from langchain.agents import AgentExecutor,create_tool_calling_agent
from tool.neo4j_tool import Neo4jTool
from tool.optimized_chroma_tool import ChromaTool
from tool.get_datetime_now import GetTimeTool
from model.my_chat_model import ChatModel
from langchain_core.prompts import ChatPromptTemplate
from langchain.memory import ConversationBufferWindowMemory
from langchain_core.callbacks import BaseCallbackHandler
# 全局变量：存储对话记忆（适用于单用户）
memory = ConversationBufferWindowMemory(k=5, return_messages=True)
# 工具只需创建一次
tools = [Neo4jTool(),ChromaTool(),GetTimeTool()]
#创建一个聊天的智能体
# 自定义回调处理器：捕获 LLM 输出的每个 token
class StreamingCallbackHandler(BaseCallbackHandler):
    def __init__(self, on_token_callback):
        self.on_token_callback = on_token_callback

    def on_llm_new_token(self, token: str, **kwargs):
        self.on_token_callback(token)
def create_agent(question,on_token):
    #获取大模型
    chat=ChatModel()
    llm=chat.get_line_model()
    #创建提示词模板
    prompt=ChatPromptTemplate([
        ("system", """
             你是一个智能的AI查询助手，你可以使用以下几个工具
                 1 get_neo4j_tool 主要用于查询知识图谱的数据，例如股东，高管，公司
                 2 get_chroma_tool 主要用于查询chroma数据库的企业背景数据
                 3 get_time_tool 主要用于查询当前时刻的时间
                 如果用户的提问的问题是关于股东、高管、公司之间的持股、投资、关联、任职、受益以及风险等问题时，请执行工具 get_neo4j_tool；
                 如果用户提问的问题是有关企业背景的调查问题，请执行工具 get_chroma_tool；请给用户准确信息，其他文本不需要
                 如果用户提问的问题是有关当前时刻时间的获取，请执行工具get_time_tool
            """),
        ("placeholder", "{history}"),  # 必须有
        ("human","{input}"),
        ("placeholder","{agent_scratchpad}") #工具调用
    ])
    #创建工具调用agent
    agent=create_tool_calling_agent(llm,tools,prompt)
    #创建智能体执行器
    agent_executor=AgentExecutor(
        agent=agent,
        tools=tools,
        memory=memory,
        verbose=True,
        handle_parsing_errors=True  # 避免因解析错误崩溃
    )
    # 使用自定义回调
    callbacks = [StreamingCallbackHandler(on_token)]
    try:
        # 启动流式执行
        response = ""
        for chunk in agent_executor.stream({"input": question}, config={"callbacks": callbacks}):
            # 每个 chunk 是 dict，可能是 'steps'、'actions'、'output'
            if "output" in chunk:
                token = chunk["output"]
                yield token  # 返回每个输出 token
                response += token
    except Exception as e:
        error_msg = f"❌ 执行出错: {str(e)}"
        yield error_msg

if __name__ == '__main__':
    create_agent("请你查询一下公司4的背景信息")
